package com.aiwiown.snackmq.api.impl;

import com.aiwiown.snackmq.api.LocalTransactionExecutor;
import com.aiwiown.snackmq.api.TransactionSendResult;
import com.aiwiown.snackmq.api.TransactionState;
import com.aiwiown.snackmq.api.TransactionStateListener;
import com.aiwiown.snackmq.api.TransactionalProducer;
import com.aiwiown.snackmq.api.config.ProducerConfig;
import com.aiwiown.snackmq.api.exception.SnackMQClientException;
import com.aiwiown.snackmq.common.message.Message;
import com.aiwiown.snackmq.common.message.MessageStatus;
import com.aiwiown.snackmq.common.message.MessageType;
import com.aiwiown.snackmq.network.client.NettyClient;
import com.aiwiown.snackmq.network.handler.ClientMessageHandler;
import lombok.extern.slf4j.Slf4j;

import java.util.UUID;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.TimeUnit;

/**
 * 事务消息生产者的默认实现。
 * <p>
 * 【已重构】: 此版本已适配父类 DefaultProducer 的连接池架构。
 * </p>
 */
@Slf4j
public class DefaultTransactionalProducer extends DefaultProducer implements TransactionalProducer, ClientMessageHandler {

    /**
     * 为每个生产者实例生成一个唯一ID，用于向 Broker 注册。
     */
    private final String producerId = UUID.randomUUID().toString();

    /**
     * 持有业务方实现的事务状态回查监听器。
     */
    private TransactionStateListener transactionStateListener;

    public DefaultTransactionalProducer(String brokerAddress, ProducerConfig producerConfig) {
        super(brokerAddress, producerConfig);
    }

    /**
     * 【核心修复】: 重写父类方法，确保在创建任何到 Broker 的新连接时，
     * 都能将本实例注册为事务回查消息的处理器。
     * 这是适配连接池架构的关键。
     */
    @Override
    protected NettyClient getOrCreateClient(String brokerAddress) throws SnackMQClientException {
        NettyClient client = super.getOrCreateClient(brokerAddress);
        // 确保此 handler 被注册到 client 的 dispatcher 中
        client.getDispatcher().registerHandler(this);
        return client;
    }

    /**
     * 设置事务状态回查监听器。
     * <p>
     * 业务方必须调用此方法并提供一个监听器实现，才能正确处理 Broker 的事务状态回查。
     * </p>
     *
     * @param listener 业务方实现的监听器
     */
    @Override
    public void setTransactionStateListener(TransactionStateListener listener) {
        this.transactionStateListener = listener;
    }

    @Override
    public void start() throws SnackMQClientException {
        super.start();
        // 在客户端启动后，立即向 Broker 注册自己，以便 Broker 知道向谁回查。
        registerSelfToBroker();
    }

    /**
     * 向 Broker 发送注册请求。
     */
    private void registerSelfToBroker() throws SnackMQClientException {
        log.info("Registering producer with ID '{}' to broker for transaction check-backs.", this.producerId);
        Message registerRequest = Message.builder()
                .type(MessageType.REGISTER_PRODUCER_REQUEST)
                .property("producerId", this.producerId)
                .build();
        // 【修复】调用父类中新的、能处理重定向的同步方法
        executeWithRedirectRetry(registerRequest, "Register Producer");
    }

    @Override
    public TransactionSendResult sendMessageInTransaction(Message message, LocalTransactionExecutor executor, Object arg) throws SnackMQClientException {
        // 1. 发送半消息 (Prepare)
        // 在发送半消息时，附带上自己的 producerId
        Message prepareRequest = message.toBuilder()
                .type(MessageType.PREPARE_TRANSACTIONAL_MESSAGE)
                .property("producerId", this.producerId)
                .build();

        // 【修复】调用父类中新的、能将消息路由到正确分区 Leader 的异步方法
        CompletableFuture<Message> prepareFuture = super.sendRequestToLeader(prepareRequest, message.getTopic(), message.getKey());

        try {
            Message prepareResponse = prepareFuture.get(DEFAULT_TIMEOUT_MS, TimeUnit.MILLISECONDS);
            if (prepareResponse.getStatus() != MessageStatus.SUCCESS) {
                throw new SnackMQClientException("Failed to send prepare message: " + prepareResponse.getErrorMessage());
            }

            String transactionId = prepareResponse.getProperty("transactionId");
            if (transactionId == null || transactionId.isEmpty()) {
                throw new SnackMQClientException("Broker did not return a transactionId for the prepare message.");
            }

            // 2. 执行本地事务
            TransactionState localTransactionState;
            try {
                localTransactionState = executor.executeLocalTransaction(message, arg);
            } catch (Exception e) {
                log.error("Exception during local transaction execution for transactionId: {}. The transaction will be rolled back.", transactionId, e);
                // 本地事务执行异常，视为需要回滚
                localTransactionState = TransactionState.ROLLED_BACK;
            }

            // 3. 发送二次确认 (Commit/Rollback)
            Message finalizationRequest;
            if (localTransactionState == TransactionState.COMMITTED) {
                finalizationRequest = Message.builder()
                        .type(MessageType.COMMIT_TRANSACTIONAL_MESSAGE)
                        .property("transactionId", transactionId)
                        .build();
            } else if (localTransactionState == TransactionState.ROLLED_BACK) {
                finalizationRequest = Message.builder()
                        .type(MessageType.ROLLBACK_TRANSACTIONAL_MESSAGE)
                        .property("transactionId", transactionId)
                        .build();
            } else { // UNKNOWN
                log.warn("Local transaction for transactionId '{}' returned UNKNOWN state. The transaction will be resolved later by the broker's check-back mechanism.", transactionId);
                // 对于 UNKNOWN 状态，我们不发送任何二次确认，等待 Broker 回查
                TransactionSendResult result = new TransactionSendResult();
                result.setStatus(MessageStatus.UNKNOWN); // 使用一个特殊状态表示
                result.setTransactionId(transactionId);
                return result;
            }

            // 【修复】调用父类中新的、能处理重定向的同步方法发送二次确认
            Message finalizationResponse = executeWithRedirectRetry(finalizationRequest, localTransactionState.toString());

            if (finalizationResponse.getStatus() != MessageStatus.SUCCESS) {
                // 这是一个严重问题，本地事务已完成，但二次确认失败。
                // 依赖 Broker 的回查机制来最终解决。
                throw new SnackMQClientException(
                        String.format("Failed to %s transaction %s. The transaction may be in an unknown state. Reason: %s",
                                localTransactionState, transactionId, finalizationResponse.getErrorMessage())
                );
            }

            // 构建最终的成功结果
            TransactionSendResult result = new TransactionSendResult();
            result.setStatus(MessageStatus.SUCCESS);
            result.setTransactionId(transactionId);
            return result;

        } catch (Exception e) {
            if (e instanceof InterruptedException) {
                Thread.currentThread().interrupt();
            }
            throw new SnackMQClientException("Transactional message sending failed.", e);
        }
    }

    // --- 实现 ClientMessageHandler 的方法，用于处理 Broker 的回查请求 ---

    @Override
    public MessageType getMessageType() {
        // 告诉分发器，我能处理这种类型的消息
        return MessageType.CHECK_TRANSACTION_STATE;
    }

    @Override
    public void handle(Message checkRequest) {
        if (this.transactionStateListener == null) {
            log.warn("Received a transaction state check request, but no TransactionStateListener is registered. Ignoring. Request: {}", checkRequest);
            return;
        }

        String transactionId = checkRequest.getProperty("transactionId");
        log.info("Received transaction state check request from broker for transactionId: {}", transactionId);

        try {
            // Broker 在回查时，会将原始的半消息放在请求的 body 中
            Message halfMessage = checkRequest.getMessages().get(0);
            // 调用业务方的逻辑来检查本地事务的最终状态
            TransactionState finalState = this.transactionStateListener.checkLocalTransactionState(halfMessage);

            Message finalizationRequest;
            if (finalState == TransactionState.COMMITTED) {
                finalizationRequest = Message.builder()
                        .type(MessageType.COMMIT_TRANSACTIONAL_MESSAGE)
                        .property("transactionId", transactionId)
                        .build();
            } else if (finalState == TransactionState.ROLLED_BACK) {
                finalizationRequest = Message.builder()
                        .type(MessageType.ROLLBACK_TRANSACTIONAL_MESSAGE)
                        .property("transactionId", transactionId)
                        .build();
            } else { // UNKNOWN
                log.info("Transaction state check for {} returned UNKNOWN. No action will be taken this time.", transactionId);
                return; // 不做任何事，等待 Broker 下一次的回查
            }

            // 【修复】将检查结果（Commit/Rollback）发送回 Broker
            executeWithRedirectRetry(finalizationRequest, "Transaction Check-back Response");
            log.info("Sent {} for transaction {} as a result of check-back.", finalizationRequest.getType(), transactionId);

        } catch (Exception e) {
            log.error("Error during handling transaction state check request for transactionId: {}", transactionId, e);
        }
    }
}